/* 
 * File:   MPLSH.cpp
 * Author: liuyi
 * 
 * Created on September 10, 2009, 9:47 PM
 */

#include "MPLSH.h"
#include "Toolkit.h"
#include "Sample.h"
#include "SampleSpace.h"
#include "KNNTraversal.h"

#include <cfloat>
#include <fstream>
#include <set>
#include <cstring>
#include <list>
#include <algorithm>

class HashFunc {
public:

    int hashUnit(float *data);

    float calcPerturb(float *data);

    void store2Stream(ofstream &fout);
    static HashFunc * loadFromStream(ifstream &fin);

    HashFunc(float *a, float b, float W, int dim) {
        m_a = new float[dim];
        memcpy(m_a, a, dim*sizeof(float));
        m_b = b;
        m_W = W;
        m_dim = dim;
    }
    ~HashFunc() {
        delete []m_a;
    }

private:    
    float *m_a;

    float m_b;

    float m_W;
    
    int m_dim;

    HashFunc() {
    }

};

int HashFunc::hashUnit(float *data) {
    float accu = 0;
    for(int i = 0; i < m_dim; i++) {
        accu += m_a[i] * data[i];
    }

    return (int) floor((accu + m_b)/m_W);
}

float HashFunc::calcPerturb(float* data) {
    float accu = 0;
    for(int i = 0; i < m_dim; i++) {
        accu += m_a[i] * data[i];
    }

    return (accu + m_b) - floor((accu + m_b) / m_W) * m_W;
}


void HashFunc::store2Stream(ofstream &fout) {
    fout << m_W << " " << m_b << " " << m_dim << " ";
    for(int i = 0; i < m_dim; i++) {
        fout << m_a[i] << " ";
    }
}

HashFunc * HashFunc::loadFromStream(ifstream &fin) {
    HashFunc *hf = new HashFunc();
    fin >> hf->m_W >> hf->m_b >> hf->m_dim;

    int dim = hf->m_dim;
    hf->m_a = new float[dim];
    for(int i = 0; i < dim; i++) {
        fin >> hf->m_a[i];
    }

    return hf;
}


class HeapElem {
public:
    vector<int> v;
    float score;

    HeapElem() {
    }

    void store2Stream(ofstream &fout);
    static HeapElem * loadFromStream(ifstream &fin);

};

void HeapElem::store2Stream(ofstream &fout) {
    int size = v.size();
    fout << size << " ";
    for(int i = 0; i < size; i++) {
        fout << v.at(i) << " ";
    }

    fout << score << " ";
}

HeapElem *HeapElem::loadFromStream(ifstream &fin) {
    HeapElem *elem = new HeapElem;

    int size, value;
    fin >> size;
    for(int i = 0; i < size; i++) {
        fin >> value;
        elem->v.push_back(value);
    }

    fin >> elem->score;

    return elem;
}

typedef struct {
public:
    int index;
    int delta;
    float value;
} Perturb;

class HashTable {
public:
    void insert(Sample *sample);

    /**
     * @return the total number of samples searched
     */
    void retreive(set<Sample *> &res, const Sample &sample, vector<HeapElem *> &pertSetsOrder);

    void store2Stream(ofstream &fout);
    static HashTable * loadFromStream(ifstream &fin);

    HashTable(int M, float *a, float *b, float W, int dim);

    ~HashTable();

private:
    class SlotElem {
    public:
        uint32_t hashCode;
        Sample *sample;

        SlotElem(uint32_t hashCode, Sample *sample): hashCode(hashCode), sample(sample){
        }

    };

    // TODO: need to set a proper value
    // static const int s_SLOT_NUM = 92083;
    static const int s_SLOT_NUM = 10000;

    int m_M;

    float m_W;

    HashFunc **m_hashFunc;

    vector<SlotElem *> m_table[s_SLOT_NUM];

    uint32_t hash(float *data);

    int mode(uint32_t code32) {
        int code = code32 % s_SLOT_NUM;
        return (code < 0 ? code + s_SLOT_NUM : code);
    }

    /**
     * @return the number of samples searched
     */
    void retreiveByCode(set<Sample *> &res, uint32_t code32);

    void genPertbSets(float *data, vector<HeapElem *> &pertSetsOrder, vector<Perturb *> *pertbSets, Perturb *pertb);

    static bool compPertb(const Perturb &a, const Perturb &b) {
        return a.value < b.value;
    }

    uint32_t hashIntArray(int *data, int size);

    HashTable() {
    }

};

//TODO: need a better integer array hash function
// string hash in PHP
uint32_t HashTable::hashIntArray(int *data, int size) {
    uint32_t code = 0, temp;
    for (int i = 0; i < size; i++) {
        code = (code << 4) + data[i];

        if ((temp = (code & 0xF0000000))) {
            code = code ^ (temp >> 24);
            code = code ^ temp;
        }
    }

    return code;
}

void HashTable::store2Stream(ofstream &fout) {
    fout << m_M << " ";
    for(int i = 0; i < m_M; i++) {
        m_hashFunc[i]->store2Stream(fout);
    }
}

HashTable * HashTable::loadFromStream(ifstream &fin) {
    HashTable *ht = new HashTable();
    fin >> ht->m_M;

    int M = ht->m_M;
    ht->m_hashFunc = new HashFunc*[M];
    for(int i = 0; i < M; i++) {
        ht->m_hashFunc[i] = HashFunc::loadFromStream(fin);
    }

    return ht;
}

HashTable::HashTable(int M, float *a, float *b, float W, int dim) {
    m_M = M;
    m_W = W;

    m_hashFunc = new HashFunc*[M];
    for(int i = 0; i < M; i++) {
        m_hashFunc[i] = new HashFunc(a + i*dim, b[i], W, dim);
    }
}

HashTable::~HashTable() {
    for(int i = 0; i < s_SLOT_NUM; i++) {
        int size = m_table[i].size();
        for(int j = 0; j < size; j++) {
            delete m_table[i].at(j);
        }
    }

    for(int i = 0; i < m_M; i++) {
        delete m_hashFunc[i];
    }
    delete []m_hashFunc;
}

uint32_t HashTable::hash(float* data) {
    int *unit = new int[m_M];
    for(int i = 0; i < m_M; i++) {
        unit[i] = m_hashFunc[i]->hashUnit(data);
    }

    uint32_t code = hashIntArray(unit, m_M);
    delete []unit;

    return code;
}

void HashTable::insert(Sample *sample) {
    uint32_t code32 = hash(sample->getFeature());
    int code = mode(code32);

    m_table[code].push_back(new SlotElem(code32, sample));
}

void HashTable::retreive(set<Sample *> &res, const Sample &sample, vector<HeapElem *> &pertbSetsOrder) {
    float *feature = sample.getFeature();

    int *unit = new int[m_M];
    for(int i = 0; i < m_M; i++) {
        unit[i] = m_hashFunc[i]->hashUnit(feature);
    }

    uint32_t code32 = hashIntArray(unit, m_M);
    retreiveByCode(res, code32);

    int count = pertbSetsOrder.size();
    vector<Perturb *> *pertbSets = new vector<Perturb *>[count];
    const int pertbNum = 2*m_M;
    Perturb *pertb = new Perturb[pertbNum];

    genPertbSets(feature, pertbSetsOrder, pertbSets, pertb);

    int *tempUnit = new int[m_M];
    int size;
    for(int i = 0; i < count; i++) {
        memcpy(tempUnit, unit, m_M*sizeof(int));

        vector<Perturb *> &v = pertbSets[i];
        size = v.size();
        for(int j = 0; j < size; j++) {
            tempUnit[v.at(j)->index] += v.at(j)->delta;
        }

        retreiveByCode(res, hashIntArray(tempUnit, m_M));
    }

    delete []tempUnit;
    delete []unit;

    delete []pertb;
    delete []pertbSets;
}

void HashTable::retreiveByCode(set<Sample *> &res, uint32_t code32) {
    int code = mode(code32);

    vector<SlotElem *>::iterator itr = m_table[code].begin();
    vector<SlotElem *>::iterator end = m_table[code].end();
    for( ; itr != end; itr++) {
        if(code32 == (*itr)->hashCode) {
            res.insert((*itr)->sample);
        }
    }
}

void HashTable::genPertbSets(float *data, vector<HeapElem *> &pertbSetsOrder, vector<Perturb *> *pertbSets, Perturb *pertb) {
    const int pertbNum = 2*m_M;

    for (int i = 0, count = 0; i < pertbNum; i += 2, count++) {
        Perturb &p1 = pertb[i];
        p1.index = count;
        p1.delta = -1;
        p1.value = m_hashFunc[count]->calcPerturb(data);

        Perturb &p2 = pertb[i + 1];
        p2.index = count;
        p2.delta = 1;
        p2.value = m_W - p1.value;
    }

    sort(pertb, pertb + pertbNum, compPertb);

    int count = pertbSetsOrder.size();
    for (int i = 0; i < count; i++) {
        vector<int> &v = pertbSetsOrder.at(i)->v;

        int size = v.size();
        for (int j = 0; j < size; j++) {
            pertbSets[i].push_back(pertb + v.at(j));
        }
    }
}


class HashModel {
public:
    void insert(Sample *sample);

    void retrieveNeighbors(vector<Sample *> &neighbors, const Sample &sample);

    /**
     * @return the total number of samples searched
     */
    long classify(int K, Sample &sample);

    void store2Stream(ofstream &fout);
    static HashModel * loadFromStream(SampleSpace &sampleSpace, vector<Sample *> &sampleList, ifstream &fin);

    HashModel(SampleSpace &sampleSpace, int M, int L, float **a, float **b, float W);

    ~HashModel();

private:
    SampleSpace &m_sampleSpace;

    int m_M;

    int m_L;

    float m_W;

    HashTable **m_hashTable;

    vector<HeapElem *> m_pertbSetsOrder;

    void initPertSetsOrder();

    bool shift(HeapElem &dst, HeapElem &src, float *Ez2);
    bool expand(HeapElem &dst, HeapElem &src, float *Ez2);

    static bool compHeapElem(const HeapElem *a, const HeapElem *b) {
        return a->score < b->score;
    }

    HashModel(SampleSpace &sampleSpace): m_sampleSpace(sampleSpace) {
    }

    HashModel(const HashModel&);

};

void HashModel::store2Stream(ofstream &fout) {
    fout << m_M << " " << m_L << " ";
    for(int i = 0; i < m_L; i++) {
        m_hashTable[i]->store2Stream(fout);
    }

    int size = m_pertbSetsOrder.size();
    fout << size << " ";
    for(int i = 0; i < size; i++) {
        m_pertbSetsOrder.at(i)->store2Stream(fout);
    }
    fout << endl;
}

HashModel * HashModel::loadFromStream(SampleSpace &sampleSpace, vector<Sample *> &sampleList, ifstream &fin) {
    HashModel *hm = new HashModel(sampleSpace);

    fin >> hm->m_M >> hm->m_L;

    int L = hm->m_L;
    hm->m_hashTable = new HashTable*[L];
    for(int i = 0; i < L; i++) {
        hm->m_hashTable[i] = HashTable::loadFromStream(fin);
    }

    int size;
    fin >> size;
    for(int i = 0; i < size; i++) {
        hm->m_pertbSetsOrder.push_back(HeapElem::loadFromStream(fin));
    }

    size = sampleList.size();
    for(int i = 0; i < size; i++) {
        hm->insert(sampleList.at(i));
    }

    return hm;
}

HashModel::HashModel(SampleSpace &sampleSpace, int M, int L, float **a, float **b, float W) : m_sampleSpace(sampleSpace) {
    m_M = M;
    m_L = L;
    m_W = W;

    m_hashTable = new HashTable*[L];
    for(int i = 0; i < L; i++) {
        m_hashTable[i] = new HashTable(M, a[i], b[i], W, m_sampleSpace.getDimension());
    }

    initPertSetsOrder();
}

HashModel::~HashModel() {
    for(int i = 0; i < m_L; i++) {
        delete m_hashTable[i];
    }

    delete []m_hashTable;

    int size = m_pertbSetsOrder.size();
    for(int i = 0; i < size; i++) {
        delete m_pertbSetsOrder.at(i);
    }
}

void HashModel::insert(Sample *sample) {
    for(int i = 0; i < m_L; i++) {
        m_hashTable[i]->insert(sample);
    }
}

void HashModel::initPertSetsOrder() {
    const int pertbNum = 2*m_M;

    float *Ez2 = new float[pertbNum];
    for (int i = 0; i < m_M; i++) {
        Ez2[i] = (i + 1)*(i + 2) * m_W * m_W / (4 * (m_M + 1)*(m_M + 2));
    }
    for (int i = m_M; i < pertbNum; i++) {
        Ez2[i] = m_W*m_W * (1 - (2*m_M - i)*1.0f / (m_M + 1) + (2*m_M - i)*(2*m_M + 1 - i)*1.0f / (4*(m_M + 1)*(m_M + 2)));
    }

    vector<HeapElem *> heap;
    HeapElem *Ai = NULL, *As = NULL, *Ae = NULL;

    Ai = new HeapElem;
    Ai->v.push_back(0);
    Ai->score = Ez2[0];
    heap.push_back(Ai);

    // TODO: alternative is 3 or 4
    const size_t NStep = 3;
    while(heap.size() > 0) {
        pop_heap(heap.begin(), heap.end(), compHeapElem);

        Ai = heap.back();
        heap.pop_back();

        if(Ai->v.size() > NStep) {
            delete Ai;

            int size = heap.size();
            for(int i = 0; i < size; i++) {
                delete heap.at(i);
            }

            break;
        }

        As = new HeapElem;
        if (shift(*As, *Ai, Ez2)) {
            heap.push_back(As);
        } else {
            delete As;
        }

        Ae = new HeapElem;
        if (expand(*Ae, *Ai, Ez2)) {
            heap.push_back(Ae);
        } else {
            delete Ae;
        }

        m_pertbSetsOrder.push_back(Ai);
    }

    delete[] Ez2;
}

bool HashModel::shift(HeapElem &dst, HeapElem &src, float *Ez2) {
    int next = src.v.back() + 1;

    int max = 2 * m_M;
    if(next >= max) {
        return false;
    }

    int size = src.v.size() - 1;
    int num;
    for(int i = 0; i < size; i++) {
        num = src.v.at(i);

        if (next == max - num - 1) {
            return false;
        }
    }

    dst.v.assign(src.v.begin(), src.v.end() - 1);
    dst.v.push_back(next);

    dst.score = src.score - Ez2[next - 1] + Ez2[next];

    return true;
}

bool HashModel::expand(HeapElem &dst, HeapElem &src, float *Ez2) {
    int next = src.v.back() + 1;

    int max = 2 * m_M;
    if(next >= max) {
        return false;
    }

    int size = src.v.size();
    int num;
    for(int i = 0; i < size; i++) {
        num = src.v.at(i);

        if (next == max - num - 1) {
            return false;
        }
    }

    dst.v.assign(src.v.begin(), src.v.end());
    dst.v.push_back(next);

    dst.score = src.score + Ez2[next];

    return true;
}

void HashModel::retrieveNeighbors(vector<Sample *> &neighbors, const Sample &sample) {
	set<Sample *> res;
	for (int i = 0; i < m_L; i++) {
		m_hashTable[i]->retreive(res, sample, m_pertbSetsOrder);
	}

	neighbors.assign(res.begin(), res.end());
}

long HashModel::classify(int K, Sample &sample) {
	vector<Sample *> sampleList;
	retrieveNeighbors(sampleList, sample);

    if(sampleList.size() == 0) {
        sample.setClass(-1);
        return 0;
    }

    KNNTraversal knnTraversal(m_sampleSpace, false);
    knnTraversal.train(sampleList);
    knnTraversal.preClassify(K);
    knnTraversal.classify(sample);

    return sampleList.size();
}

void MPLSH::train(vector<Sample *> &sampleList) {
    KNN::train(sampleList);

    int size = m_sampleList->size();
    for(int i = 0; i < size; i++) {
        m_hashModel->insert(m_sampleList->at(i));
    }
}

void MPLSH::store2File(const char *modelPath) {
    ofstream fout(modelPath);

    m_sampleSpace->store2Stream(fout);
    m_sampleSpace->storeSample2Stream(fout, *m_sampleList);
    m_hashModel->store2Stream(fout);

    fout.close();
}

MPLSH * MPLSH::loadFromFile(const char *modelPath) throw(FileNotFoundException) {
    ifstream fin(modelPath);
    if(!fin) {
        string filePath(modelPath);
        throw FileNotFoundException(filePath + " does not exits");
    }

    SampleSpace *ss = SampleSpace::loadFromStream(fin);

    MPLSH *mplsh = new MPLSH(*ss);

    mplsh->m_sampleList = new vector<Sample *>;
    ss->loadSampleFromStream(fin, *mplsh->m_sampleList);

    mplsh->m_hashModel = HashModel::loadFromStream(mplsh->getSampleSpace(), *mplsh->m_sampleList, fin);

    fin.close();

    delete ss;

    return mplsh;
}

void MPLSH::retrieveNeighbors(vector<Sample *> &neighbors, const Sample &sample) {
	m_hashModel->retrieveNeighbors(neighbors, sample);
}

void MPLSH::classify(Sample &sample) {
    m_hashModel->classify(m_K, sample);
}

float MPLSH::classifyWithSearchRatio(Sample &sample) {
    long searchedNum = m_hashModel->classify(m_K, sample);
    return searchedNum * 1.0f / m_sampleList->size();
}

void MPLSH::preTrain(int M, int L, float** a, float** b, float W) {
    releaseModel();
    m_hashModel = new HashModel(*m_sampleSpace, M, L, a, b, W);
}

MPLSH::MPLSH(SampleSpace &sampleSpace, bool copyList) : KNN(sampleSpace, copyList) {
    m_hashModel = NULL;
}

void MPLSH::releaseModel() {
    releaseSampleList();

    if(m_hashModel != NULL) {
        delete m_hashModel;
        m_hashModel = NULL;
    }
}
